#include <math.h>
#include <float.h>

#include "trilinear_kernel.h"

__global__ void TriLinearForward(const float* lut, const float* image, float* output, const int dim, const int shift, \
                                 const float binsize, const int width, const int height, const int batch, const int bit_num)
{
    unsigned int row = blockDim.y*blockIdx.y + threadIdx.y;
    unsigned int col = blockDim.x*blockIdx.x + threadIdx.x;

    const float scale = (1.0f / (float) ( (1 << bit_num) - 1) ) * (MAX_LEVEL - 1);

    int index = row * width + col;
    int index_r = 3 * index + 0;
    int index_g = 3 * index + 1;
    int index_b = 3 * index + 2;

    float r = image[index_r] * scale;
    float g = image[index_g] * scale;
    float b = image[index_b] * scale;

    int r_id = int(r * (1 << MAX_SHIFT_NUM)) >> MAX_SHIFT_NUM;
    int g_id = int(g * (1 << MAX_SHIFT_NUM)) >> MAX_SHIFT_NUM;
    int b_id = int(b * (1 << MAX_SHIFT_NUM)) >> MAX_SHIFT_NUM;

    float r_d = (r - r_id);
    float g_d = (g - g_id);
    float b_d = (b - b_id);

    int id000 = r_id + g_id * dim + b_id * dim * dim;
    int id100 = r_id + 1 + g_id * dim + b_id * dim * dim;
    int id010 = r_id + (g_id + 1) * dim + b_id * dim * dim;
    int id110 = r_id + 1 + (g_id + 1) * dim + b_id * dim * dim;
    int id001 = r_id + g_id * dim + (b_id + 1) * dim * dim;
    int id101 = r_id + 1 + g_id * dim + (b_id + 1) * dim * dim;
    int id011 = r_id + (g_id + 1) * dim + (b_id + 1) * dim * dim;
    int id111 = r_id + 1 + (g_id + 1) * dim + (b_id + 1) * dim * dim;

    float w000 = (1-r_d)*(1-g_d)*(1-b_d);
    float w100 = r_d*(1-g_d)*(1-b_d);
    float w010 = (1-r_d)*g_d*(1-b_d);
    float w110 = r_d*g_d*(1-b_d);
    float w001 = (1-r_d)*(1-g_d)*b_d;
    float w101 = r_d*(1-g_d)*b_d;
    float w011 = (1-r_d)*g_d*b_d;
    float w111 = r_d*g_d*b_d;

    float out = w000 * lut[id000] + w100 * lut[id100] +
                w010 * lut[id010] + w110 * lut[id110] +
                w001 * lut[id001] + w101 * lut[id101] +
                w011 * lut[id011] + w111 * lut[id111];
    output[index_r] = out * ( (1 << bit_num) - 1);

    out = w000 * lut[id000 + shift] + w100 * lut[id100 + shift] +
          w010 * lut[id010 + shift] + w110 * lut[id110 + shift] +
          w001 * lut[id001 + shift] + w101 * lut[id101 + shift] +
          w011 * lut[id011 + shift] + w111 * lut[id111 + shift];
    output[index_g] = out * ( (1 << bit_num) - 1);

    out = w000 * lut[id000 + shift * 2] + w100 * lut[id100 + shift * 2] +
          w010 * lut[id010 + shift * 2] + w110 * lut[id110 + shift * 2] +
          w001 * lut[id001 + shift * 2] + w101 * lut[id101 + shift * 2] +
          w011 * lut[id011 + shift * 2] + w111 * lut[id111 + shift * 2];
    output[index_b] = out * ( (1 << bit_num) - 1);
}

int TriLinearForwardGpu(const float* lut, const float* image, float* output, const int lut_dim, const int shift, \
                        const float binsize, const int width, const int height, const int batch, const int bit_num, \
                        cudaStream_t stream)
{
    dim3 dimBlock(16, 16);
    dim3 dimGrid((width + dimBlock.x - 1) / (dimBlock.x), (height + dimBlock.y - 1) / (dimBlock.y));
    TriLinearForward<<<dimGrid, dimBlock, 0, stream>>>(lut, image, output, lut_dim, shift, binsize, width, height, batch, bit_num);

    return 0;
}

